from scipy.misc import imread
from scipy.misc import imsave
import numpy as np
import logging
from logging import DEBUG as debug
from logging import log
from math import floor
from scipy.ndimage.measurements import center_of_mass
kernel = np.array(
    [
        [0, 1, 0],
        [1, 1, 1],
        [0, 1, 0]
    ]
)

kernel2 = np.array([
    [1, 1, 1],
    [1, 1, 1],
    [1, 1, 1]
])


def dilate(img, k):
    assert len(img.shape) == 2
    assert k.shape[0] == k.shape[1]
    sizeKernal = floor(k.shape[0] / 2)
    out = img > 0
    out = np.double(out)
    for row in range(sizeKernal, img.shape[0] - sizeKernal):
        for col in range(sizeKernal, img.shape[1] - sizeKernal):
            if img[row, col]:
                out[
                (row - sizeKernal):(row + sizeKernal + 1),
                (col - sizeKernal):(col + sizeKernal + 1)
                ] += k
    return out > 0


def erosion(img, k):
    invImg = img <= 0
    outImg = dilate(invImg, k)
    return outImg <= 0


def openOp(img, k):
    return dilate(erosion(img, k), k)


def closeOp(img, k):
    return erosion(dilate(img, k), k)


def boundary(img, k):
    img = img > 0
    return img - erosion(img, k)


def filling(img, k, iPoint):
    invImg = img <= 0
    img = np.zeros(invImg.shape)
    img[iPoint[0], iPoint[1]] = 1
    while (1):
        prevImg = img
        img = dilate(img, k) * invImg
        comp = np.abs(prevImg - img)
        # log(debug,np.sum(comp))
        if np.sum(comp) == 0:
            break
    return img


def connected(img, k, iPoint):
    img = img > 0
    X = np.zeros(img.shape)
    X[iPoint[0], iPoint[1]] = 1
    while 1:
        X_prev = X
        X = dilate(X, kernel2) * img
        comp = np.abs(X - X_prev)
        if np.sum(comp) == 0:
            break
        # log(debug, np.sum(X))
    return X


def partI():
    baseDir = 'data/'
    files = baseDir + 'noisy_fingerprint.tif'
    oriimg = imread(files)
    log(debug, oriimg.shape)
    dilateImg = dilate(oriimg, kernel)
    erosionImg = erosion(oriimg, kernel)
    openImg = openOp(oriimg, kernel)
    closeImg = closeOp(oriimg, kernel)
    imsave('out/hw8_dilate_result.jpg', dilateImg)
    imsave('out/hw8_erosion_result.jpg', erosionImg)
    imsave('out/hw8_open_result.jpg', openImg)
    imsave('out/hw8_close_result.jpg', closeImg)
    imsave('out/hw8_dilate_original.jpg', oriimg)


def partII():
    baseDir = 'data/'
    boundaryOriginal = imread(baseDir + 'licoln_from_penny.tif')
    boundaryImg = boundary(boundaryOriginal, kernel)
    imsave('out/hw8_boundary_original.jpg', boundaryOriginal)
    imsave('out/hw8_boundary_result.jpg', boundaryImg)

    fillPoint = [
        [55, 63],
        [180, 4],
        [159, 94],
        [298, 106],
        [449, 56],
        [202, 179],
        [149, 261],
        [308, 230],
        [460, 233],
        [36, 363],
        [141, 456],
        [236, 414],
        [371, 386],
        [472, 403],
        [489, 506],
        [45, 182]
    ]
    fillOriginal = imread(baseDir + 'region_filling_reflections.tif')
    log(debug, fillOriginal.shape)
    fillOriginal[0:fillOriginal.shape[0], 0] = 255
    fillOriginal[0:fillOriginal.shape[0], fillOriginal.shape[1] - 1] = 255
    #
    # fillOriginal[0, 0:fillOriginal.shape[1]] = 255
    # fillOriginal[fillOriginal.shape[1] - 1, 0:fillOriginal.shape[1]] = 255
    # fillingResult = fillOriginal
    # for p in fillPoint:
    #     fillingResult += filling(fillOriginal, kernel, p)
    # fillingResult = fillOriginal > 0
    #
    # imsave('out/hw8_fill_original.jpg', fillOriginal)
    # imsave('out/hw8_fill_result.jpg', fillingResult)

    connPoint = [
        [132,188],
        [154, 172],
        [158,195],
        [94,271],
        [91,357],
        [87,398],
        [243,415],
    ]
    connectedOriginal = imread(baseDir + 'chickenfilet_with_bones.tif')
    afterThresh = connectedOriginal > 200
    afterThresh = openOp(afterThresh, np.ones((3,3)))
    connectedResult = np.zeros(connectedOriginal.shape)
    for P in connPoint:
        X = connected(afterThresh, kernel2, P)
        connectedResult += X
        center = center_of_mass(X)
        log(debug, 'Area: %d, center(%f,%f)' % (np.sum(X), center[0], center[1]))
    connectedOriginal = connectedOriginal > 0
    imsave('out/hw8_connected_original.jpg', connectedOriginal)
    imsave('out/hw8_connected_thresh.jpg', afterThresh)
    imsave('out/hw8_connected_result.jpg', connectedResult)


def main():
    # Part I
    FORMAT = '	[%(levelname)-5s]%(asctime)-8s %(filename)s:%(lineno)d %(message)s'
    DATEFORMAT = '%H:%M:%S'
    logging.basicConfig(level=logging.DEBUG, format=FORMAT, datefmt=DATEFORMAT)
    logging.log(logging.DEBUG, "Start")

    partI()
    partII()


if __name__ == '__main__':
    main()
